# coding: utf-8
"""
Darknet 53
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from .block import *

class Darknet53(nn.Module):
    def __init__(self, ch_in=3):
        super(Darknet53, self).__init__()

        # conv1
        self.conv11 = DBL(ch_in, 32, kernel=3, stride=1, pad=1)
        self.conv12 = DBL(32, 64, kernel=3, stride=2, pad=1)

        # 1
        self.res1 = dark_res(64)
        self.res1_2 = DBL(64, 128, kernel=3, stride=2, pad=1)

        # 2
        self.res2 = nn.ModuleList([dark_res(128)] * 2)
        self.res2_2 = DBL(128, 256, kernel=3, stride=2, pad=1)

        # 8
        self.res3 = nn.ModuleList([dark_res(256)] * 8)
        self.res3_2 = DBL(256, 512, kernel=3, stride=2, pad=1)

        # 8
        self.res4 = nn.ModuleList([dark_res(512)] * 8)
        self.res4_2 = DBL(512, 1024, kernel=3, stride=2, pad=1)

        # 4
        self.res5 = nn.ModuleList([dark_res(1024)] * 4)

    def forward(self, x):
        x = self.conv11(x)
        x = self.conv12(x)

        # 1
        x = self.res1(x)
        x = self.res1_2(x)

        # 2
        for l in self.res2:
            x = l(x)
        x = self.res2_2(x)

        # 8
        for l in self.res3:
            x = l(x)
        x1 = x
        x = self.res3_2(x)

        # 8
        for l in self.res4:
            x = l(x)
        x2 = x
        x = self.res4_2(x)

        # 4
        for l in self.res5:
            x = l(x)

        return x1, x2, x

if __name__ == "__main__":
    x = torch.randn(1, 3, 416, 416)
    net = Darknet53(3)
    y = net(x)
    print("{} {} {}".format(y[0].shape, y[1].shape, y[2].shape))
